python3 main_pretrain.py \
    --dataset $1 \
    --backbone resnet18 \
    --data_dir ./datasets \
    --max_epochs $2 \
    --devices 0 \
    --accelerator gpu \
    --precision 16 \
    --optimizer sgd \
    --scheduler warmup_cosine \
    --lr 0.5 \
    --classifier_lr 0.1 \
    --weight_decay 1e-5 \
    --batch_size 256 \
    --num_workers 4 \
    --crop_size 32 \
    --brightness 0.4 \
    --contrast 0.4 \
    --saturation 0.4 \
    --hue 0.1 \
    --gaussian_prob 0.0 0.0 \
    --crop_size 32 \
    --num_crops_per_aug 1 1 \
    --zero_init_residual \
    --name rdmsimsiam-$1-$2ep-$3_$5 \
    --project solo-learn \
    --entity unitn-mhug \ # Your entity 
    --wandb \
    --save_checkpoint \
    --auto_resume \
    --method rdmsimsiam \
    --proj_hidden_dim 2048 \
    --proj_output_dim 2048 \
    --pred_type $3 \
    --pred_location $4 \
    --sigma $5 \

